import pandas as pd
import itertools
import seaborn as sns
import matplotlib.pyplot as plt
import re

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
data = pd.read_csv('data/covid_drug_papers_sample.csv')

drug_terms = ['chloroquine', 'hydroxychloroquine', 'azithromycin', 'remdesivir', 'lopinavir', 'ritonavir', 'favipiravir', 'tocilizumab', 'dexamethasone', 'zinc', 'doxycycline', 'ivermectin', 'heparin', 'baricitinib']
pattern = re.compile('(' + '|'.join(drug_terms) + ')', re.IGNORECASE)

def get_drugs(text):
    if isinstance(text, str):
        found = pattern.findall(text)
        return sorted({word.lower() for word in found})
    return []

data['drugs'] = data['abstract'].apply(get_drugs)

pairs = {}
for items in data['drugs']:
    for a, b in itertools.combinations(items, 2):
        key = tuple(sorted((a, b)))
        pairs[key] = pairs.get(key, 0) + 1

unique = sorted({word for pair in pairs for word in pair})
co = pd.DataFrame(0, index=unique, columns=unique)
for (a, b), count in pairs.items():
    co.loc[a, b] = count
    co.loc[b, a] = count

plt.figure(figsize=(8, 6))
sns.heatmap(co, cmap='Reds', annot=True, fmt='g')
plt.tight_layout()
plt.savefig('第二次平时作业_第二题_图1.png', dpi=300)
plt.close()

dosage_pattern = re.compile(r'(\\d+\\s*(?:mg|mg/kg|g|iu))', re.IGNORECASE)
rows = []
for index, row in data.iterrows():
    text = str(row['abstract'])
    for drug in row['drugs']:
        for match in re.finditer(drug, text, re.IGNORECASE):
            window = text[match.start(): match.end() + 40]
            dose = dosage_pattern.search(window)
            if dose:
                rows.append({'paper_id': row['paper_id'], 'drug': drug, 'dosage': dose.group(0)})

dosages = pd.DataFrame(rows)
dosages.to_csv('第二次平时作业_第二题_剂量.csv', index=False)
